BeginPackage["SpinBath`",{"MathMR`", "MRData`"}];
LoadLattice			::usage="Create a lattice according to the given structure.";
IsCoord				::usage="An option for LoadLattice, if true, gives the coordiantes, else gives the index.";
GenerateBath		::usage="Generate bath spins on a given lattice with a given abundance.";

GenerateSiCBath	::usage="Generate a SiC bath";
GenerateSiCBathSeries::usage="";
Coord2HSiC			::usage="coord of 2HSiC";
Coord4HSiC			::usage="coord of 4HSiC";
Coord6HSiC			::usage="coord of 6HSiC";
LattSiC			::usage="generate SiC lattice";
SiCgCoord			::usage="prepend element names to coordinates";

LookUpDist			::usage="look up distance";

LatticeNeighbor		::usage="Calculate the neighbors (coordinate shells) of the given lattice.";
BathDistAnalyze		::usage="Calculate the distance distribution (the spin pair numbers for a centain distance).";
GetCoherentCluster	::usage="";
GenerateLinkList	::usage="";
LinkList			::usage="";
ConnectAbsorb		::usage="";
CoarseGrainClusters	::usage="";
Spin2Info			::usage="Show the basic info of the 2-spin cluster.";
Spin3Info			::usage="Show the basic info of the 3-spin cluster.";
ExportBathFile		::usage="Export the generated spin bath.";

Index2Coord::usage="convert index to coordinates using the diamond lattice";
GenerateC13Bath::usage="randomly generate the C-13 bath spin coordinates for given isotope abundance";

(*=================================================================*)
(*Load lattice structures*)
Begin["`LatticePrivate`"];

Options[LoadLattice]={IsCoord->True};
LoadLattice[name_,lattConst_,range_,opt___] := 
	Module[{res,isCoord,idx,minX,maxX,minY,maxY,minZ,maxZ,len}, 
		isCoord=IsCoord/.{opt}/.Options[LoadLattice];
		Which[
			(*Diamond Lattice*)
			name=="Diamond",
				{{minX,maxX},{minY,maxY},{minZ,maxZ}}=range;
				
				len=$DiamondRList//Length;
				idx=Flatten[
					Table[(i*$e1+j*$e2+k*$e3)+$DiamondRList[[r]],
						{i,minX,maxX},{j,minY,maxY},{k,minZ,maxZ},{r,1,len}
						],3
					];
				res=If[isCoord,lattConst*idx,idx];
				,
			
			(*Error*)
			True,
			Print["Error! No lattice data."];
			Return[{}];
		];
		Return[res];
	]
	
End[];

(*=================================================================*)
(*Generate Bath*)
Begin["`BathPrivate`"];

Index2Coord[{i_, j_, k_, r_}] := ({i, j, k} + $DiamondRList[[r]])*$DiamondLatticeConst

Coord2HSiC[atom_, na_, nb_, nc_] := 
 Module[{res, layerIdx, layerDist, inplaneShift, inplane, zVect},
   layerDist = $2HSiCconstC/$2HSiCLayerNumber;
   
   layerIdx = Mod[nc, $2HSiCLayerNumber];
   inplaneShift = $2HSiCLayer[[layerIdx + 1]];
   inplane = 
    $2HSiCconstA*(na*$SiCUnitVectA + nb*$SiCUnitVectB + inplaneShift);
   
   zVect = layerDist*{0, 0, nc};
   
   res = inplane + zVect;
   If[atom == "C", res = res + $SiCShift*layerDist];
   Return[res];
   ] /; (atom == "Si" || atom == "C")

Coord4HSiC[atom_, na_, nb_, nc_] := 
 Module[{res, layerIdx, layerDist, inplaneShift, inplane, zVect},
   layerDist = $4HSiCconstC/$4HSiCLayerNumber;
   
   layerIdx = Mod[nc, $4HSiCLayerNumber];
   inplaneShift = $4HSiCLayer[[layerIdx + 1]];
   inplane = 
    $4HSiCconstA*(na*$SiCUnitVectA + nb*$SiCUnitVectB + inplaneShift);
   
   zVect = layerDist*{0, 0, nc};
   
   res = inplane + zVect;
   If[atom == "C", res = res + $SiCShift*layerDist];
   Return[res];
   ] /; (atom == "Si" || atom == "C")

Coord6HSiC[atom_, na_, nb_, nc_] := 
 Module[{res, layerIdx, layerDist, inplaneShift, inplane, zVect},
   layerDist = $6HSiCconstC/$6HSiCLayerNumber;
   
   layerIdx = Mod[nc, $6HSiCLayerNumber];
   inplaneShift = $6HSiCLayer[[layerIdx + 1]];
   inplane = 
    $6HSiCconstA*(na*$SiCUnitVectA + nb*$SiCUnitVectB + inplaneShift);
   
   zVect = layerDist*{0, 0, nc};
   
   res = inplane + zVect;
   If[atom == "C", res = res + $SiCShift*layerDist];
   Return[res];
   ] /; (atom == "Si" || atom == "C")

LattSiC[atom_, typeFunc_, maxNa_, maxNb_, maxNc_]:=
	Module[{res},
		res=Flatten[Table[typeFunc[atom, a, b, c], {a, -maxNa, maxNa}, {b, -maxNb, maxNb}, {c, -maxNc, maxNc}], 2];	
		Return[res];
	]
   
GenerateSiCBath[type_, cAbund_, siAbund_,cSeed_,siSeed_,rCut_, maxIdx_:20]:=
	Module[{res, Clattice, Silattice, cCoord, siCoord},
		Which[
			type=="2H",
			Clattice=LattSiC["C", Coord2HSiC, maxIdx, maxIdx, maxIdx];
			Silattice=LattSiC["Si", Coord2HSiC, maxIdx, maxIdx, maxIdx];
			,
			type=="4H",
			Clattice=LattSiC["C", Coord4HSiC, maxIdx, maxIdx, maxIdx];
			Silattice=LattSiC["Si", Coord4HSiC, maxIdx, maxIdx, maxIdx];
			,
			type=="6H",
			Clattice=LattSiC["C", Coord6HSiC, maxIdx, maxIdx, maxIdx];
			Silattice=LattSiC["Si", Coord6HSiC, maxIdx, maxIdx, maxIdx];
		];
		cCoord=Select[GenerateBath[Clattice, cAbund,cSeed],Norm[#]<rCut &];
		siCoord=Select[GenerateBath[Silattice, siAbund,siSeed],0.0001<Norm[#]<rCut &];
		res={cCoord, siCoord};
		
		Return[res];
	]/;MemberQ[{"2H","4H","6H"},type]
SiCgCoord[coords_]:=(Prepend[#,"C"]&/@coords[[1]])~Join~(Prepend[#,"Si"]&/@coords[[2]])

GenerateSiCBathSeries[type_, cAbund_, siAbund_,cSeedList_,siSeedList_,rCut_, maxIdx_:20]:=
	Module[{res, Clattice, Silattice, cCoord, siCoord,len},
		len= Length[cSeedList];
		Which[
			type=="2H",
			Clattice=LattSiC["C", Coord2HSiC, maxIdx, maxIdx, maxIdx];
			Silattice=LattSiC["Si", Coord2HSiC, maxIdx, maxIdx, maxIdx];
			,
			type=="4H",
			Clattice=LattSiC["C", Coord4HSiC, maxIdx, maxIdx, maxIdx];
			Silattice=LattSiC["Si", Coord4HSiC, maxIdx, maxIdx, maxIdx];
			,
			type=="6H",
			Clattice=LattSiC["C", Coord6HSiC, maxIdx, maxIdx, maxIdx];
			Silattice=LattSiC["Si", Coord6HSiC, maxIdx, maxIdx, maxIdx];
		];
		
		res=Table[
			cCoord=Select[GenerateBath[Clattice, cAbund,cSeedList[[i]] ],Norm[#]<rCut &];
			siCoord=Select[GenerateBath[Silattice, siAbund,siSeedList[[i]] ],0.0001<Norm[#]<rCut &];
			{cCoord, siCoord},
			{i,1,len}
			];
		Return[res];
	]/;MemberQ[{"2H","4H","6H"},type] && (Length[cSeedList]==Length[siSeedList])


GenerateC13Bath[c13_, seed_, rNat_: 50.0] := 
 Module[{res, rBound, nBox, siteNum, C13Num, boxIdx, inBoxIdx, idx, 
   C13Index, errPercent},
  rBound = rNat (c13/$DiamondNatureAbundance)^(-1/3);
  nBox = Ceiling[rBound/$DiamondLatticeConst];
  siteNum = (2*nBox)^3*8;
  C13Num = Floor[siteNum*c13];
  
  SeedRandom[seed];
  boxIdx = RandomInteger[{-nBox, nBox}, {C13Num, 3}];
  inBoxIdx = RandomInteger[{1, 8}, C13Num];
  idx = MapThread[Append[#1, #2] &, {boxIdx, inBoxIdx}];
  C13Index = Complement[idx, {{0, 0, 0, 0}, {0, 0, 0, 1}}];
  errPercent = 100*((C13Index // Length)/siteNum - c13)/c13;
  
  If[Abs[errPercent] > 1., Print["Error = ", errPercent]];
  
  res = Index2Coord /@ C13Index;
  Return[res];
  ]

GenerateBath[lattCoord_,abund_,seed_]:=
	Module[{res,len,rand,resUnsort},
		len=lattCoord//Length;
		SeedRandom[seed];
		rand=RandomReal[{0.,1.},len];
		resUnsort=lattCoord[[Position[rand,a_/;a<abund]//Flatten]];
		res=Sort[resUnsort,Norm[#1]<Norm[#2]&];
		Return[res];
	]

LookUpDist[i_,j_,nbath_,distList_]:=
	Module[{res},
		res=If[i>=j,distList[[i*(i-1)/2+j]],distList[[j*(j-1)/2+i]]];
		Return[res];
	]/;i<=nbath&&j<=nbath

End[];

(*=================================================================*)
(*Analyze spin bath*)
Begin["`AnalysePrivate`"];

LatticeNeighbor[name_,lattConst_,range_]:=
	Module[{res,coord,dist,neighborDist,numNeighbor,rg},
		rg={{-range,range},{-range,range},{-range,range}};
		coord=LoadLattice[name,lattConst,rg,IsCoord->False];
		dist=Norm/@coord;
		neighborDist=Sort[Select[Union[dist],#<range&],#1<#2&];
		numNeighbor=Count[dist,#]&/@neighborDist;
		res=Transpose[{neighborDist*lattConst,numNeighbor}];
		Return[res];
	]

BathDistAnalyze[linkinfo_,bathcoord_,neighborList_,eps_:10^-6]:=
	Module[{res,len,len1,distlist,unionDist,numDist},
		len=linkinfo//Length;
		distlist=linkinfo[[All,3]];
		unionDist=Union[distlist,SameTest->(Abs[#1-#2]<eps&)];
		len1=unionDist//Length;
		If[Norm[neighborList[[2;;len1+1,1]]-unionDist]>10^-10,Print["Mismatch!"]];
		numDist=Count[distlist,a_/;Abs[a-#]<eps]&/@unionDist;
		res=Transpose[{Range[len1],unionDist,numDist}];
		Return[res];
	]

GetCoherentCluster[linkinfo_,bathcoord_,neighborList_,neighborDepth_]:=
	Module[{res,len,selpair,r,tf,triPos,res2spin,res3spin},
		r=(neighborList[[neighborDepth+1,1]]+neighborList[[neighborDepth+2,1]])/2;
		selpair=Select[linkinfo,(#[[3]]<r)&][[All,{1,2}]];
		len=selpair//Length;

		tf=Outer[Length[#1\[Union]#2]==3&,selpair,selpair,1];
		triPos=Union[Sort/@Position[tf,True]];

		res2spin=Delete[selpair,List/@(triPos//Flatten)];
		res3spin=Union[selpair[[#]]//Flatten]&/@triPos;
		res={res2spin,res3spin};
		Return[res];
	]
	

	
GenerateLinkList[coord_,thresholddist_] := 
	Module[{res,lk}, 
		lk=Install[$MRMathLinkPath<>"LinkList/LinkList"];
		res=Partition[LinkList[Flatten[coord],thresholddist],2];
		Uninstall[lk];
		Return[res];
	]
		
ConnectAbsorb[absorbedList_, element_] :=
	Module[{res, len1, len2, len3, uList},
		len1 = absorbedList[[1]] // Length;
		len2 = element // Length;
		uList = Union[absorbedList[[1]], element];
		len3 = uList // Length;

		res = If[len1 + len2 > len3,
  			ReplacePart[absorbedList, 1 -> uList]
  			,
  			Append[absorbedList, element]
  			];
		Return[res];
	]

CoarseGrainClusters[lklst_,fullcoord_] := 
	Module[{res, remain, foldres},
		res = List/@Complement[Range[fullcoord // Length], Union[lklst // Flatten]];
		remain = lklst;
		While[Length[remain] > 1,
			foldres = FixedPoint[Fold[ConnectAbsorb, #[[{1}]], #[[2 ;;]]] &, remain];
			res = Append[res, foldres[[1]]];
			If[Length[remain] > 1, remain = foldres[[2 ;;]] ];
		];

		If[Length[remain] == 1, res=Append[res, remain[[1]]]];
		res=Sort[res, Length[#1] < Length[#2] &];
		Return[res];
	]

Spin2Info[spin2idx_,bathcoord_]:=
	{#,bathcoord[[#]],Norm[Differences[bathcoord[[#]]]],Norm[Mean[bathcoord[[#]]]]}&/@spin2idx
	
Spin3Info[spin3idx_,bathcoord_]:=
	{#,bathcoord[[#]],Norm[Mean[bathcoord[[#]]]]}&/@spin3idx

End[];

(*=================================================================*)
(*Export Bath*)
Begin["`ExportBath`"];

ExportBathFile[file_,coord_,elementList_]:=
	Module[{res,num,st},
		num=coord//Length;
		st=OpenWrite[file];
		Write[st,num];
		WriteString[st,"#Commet\n"];
		Do[
			WriteString[st,elementList[[i]]<>"\t"<>ToString[coord[[i,1]]]<>"\t"<>
				ToString[coord[[i,2]]]<>"\t"<>ToString[coord[[i,3]]]<>"\n"];
			,{i,1,num}
		];
		res=Close[st];
		Return[res];
	]/;Length[coord]==Length[elementList]

ExportBathFile[file_,gcoord_]:=
	Module[{res,num,st,coord,elementList},
		num=gcoord//Length;
		coord=gcoord[[All,{2,3,4}]];
		elementList=gcoord[[All,1]];
		
		st=OpenWrite[file];
		Write[st,num];
		WriteString[st,"#Commet\n"];
		Do[
			WriteString[st,elementList[[i]]<>"\t"<>ToString[coord[[i,1]]]<>"\t"<>
				ToString[coord[[i,2]]]<>"\t"<>ToString[coord[[i,3]]]<>"\n"];
			,{i,1,num}
		];
		res=Close[st];
		Return[res];
	]

End[];

(*=================================================================*)
EndPackage[]; 
